import cv2 as cv
import numpy as np
import os
import sys

BASE_DIR, FILE_NAME = os.path.split(__file__)
path = '../../../../large_data/video/traffic/traffic_20s.avi'
VIDEO_PATH = os.path.join(BASE_DIR, path)
history_cnt = 500
ALPHA = 1e-3
INTERVAL = 1000 // 25

print('Training ...')
bg_s = cv.createBackgroundSubtractorMOG2(history=history_cnt, detectShadows=False)
video = cv.VideoCapture(VIDEO_PATH)
cnt = 0
while True:
    ret, img = video.read()
    if not ret:
        break
    cnt += 1
    bg_s.apply(img, None, ALPHA)
    if cnt >= history_cnt:
        break
if cnt < history_cnt:
    print(f'Warning: not enough frames! Need {history_cnt}, actually only {cnt}!', file=sys.stderr)
video.release()
print('Trained.')

video = cv.VideoCapture(VIDEO_PATH)
cnt = 0
while True:
    ret, img = video.read()
    if not ret:
        print('Video over.')
        break
    cnt += 1
    mask = bg_s.apply(img, None, 0.)
    mask_ = mask.copy()

    kernel = np.ones((3, 3), dtype=np.uint8)
    mask = cv.morphologyEx(mask, cv.MORPH_OPEN, kernel)
    mask = cv.morphologyEx(mask, cv.MORPH_CLOSE, kernel)
    mask = cv.dilate(mask, kernel, iterations=2)

    final = np.concatenate([mask_, mask], axis=0)

    cv.imshow('mask', final)
    k = cv.waitKey(INTERVAL) & 0xFF
    if 27 == k:
        print('Interrupted by user.')
        break
print(f'Processed {cnt} frames of the video.')
video.release()

cv.destroyAllWindows()
